# src/core/template_matcher.py
import os

import cv2
import numpy as np
from abc import ABC

class TemplateMatcher(ABC):
    def __init__(self, template_dir):
        self.template_dir = template_dir
        self.templates = self._load_templates()

    def _load_templates(self):
        """
        自动加载 template_dir 目录下所有形如 '0.png', '1.png', ... 的图片
        支持任意数量模板，不限于 0-9
        """
        templates = {}

        if not os.path.exists(self.template_dir):
            raise FileNotFoundError(f"模板目录不存在: {self.template_dir}")

        # 遍历目录下所有文件
        for filename in os.listdir(self.template_dir):
            # 只处理以 .png 结尾的文件
            if filename.endswith('.png'):
                # 尝试提取文件名中的数字（如 '0.png' -> 0, 'special.png' -> 跳过）
                name = filename[:-4]  # 去掉 .png
                if name.isdigit():
                    digit = int(name)
                    path = os.path.join(self.template_dir, filename)
                    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
                    if img is not None:
                        templates[digit] = img
                    else:
                        print(f"警告：无法加载图像 {path}")

        if not templates:
            raise RuntimeError(f"未在目录中找到任何有效模板: {self.template_dir}")

        print(f"成功加载 {len(templates)} 个模板: {sorted(templates.keys())}")
        return templates

    def match_digit(self, image, threshold=0.7):
        """
        直接比较整图相似度，找到第一个 >= threshold 的匹配。
        """
        # 1. 预处理：转为灰度 NumPy 数组
        # 如果 image 已经是 numpy.ndarray 且为灰度图，可跳过转换
        if isinstance(image, np.ndarray):
            if len(image.shape) == 3:  # 彩色图
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:  # 灰度图
                gray = image
        else:  # 假设是 PIL Image
            gray = np.array(image.convert('L'))

        # 2. 遍历每个模板
        for digit, template in self.templates.items():
            # 3. 计算归一化互相关 (NCC) - 这就是 matchTemplate 在 (0,0) 位置的值
            # 使用 OpenCV 的 matchTemplate 在单点计算，或手动计算
            # 方法一：使用 OpenCV (推荐，C++ 实现，最快)
            res = cv2.matchTemplate(gray, template, cv2.TM_CCOEFF_NORMED)
            similarity = res[0, 0]  # 因为图像大小相同，只在 (0,0) 位置有值

            # 方法二：手动计算 NCC (供理解，通常比 OpenCV 慢)
            # mean_gray = np.mean(gray)
            # mean_template = np.mean(template)
            # numerator = np.sum((gray - mean_gray) * (template - mean_template))
            # denominator = np.sqrt(np.sum((gray - mean_gray)**2) * np.sum((template - mean_template)**2))
            # similarity = numerator / (denominator + 1e-8)  # 防止除零

            # 4. 检查是否达到阈值
            if similarity >= threshold:
                return digit  # ✅ 立即返回！找到匹配

        # 5. 如果所有模板都不满足阈值，返回 0
        return -1
